# -*- coding: utf-8 -*-
# time: 2025/3/27 16:56
# file: chromadb_ollama_ARG.py
# author: hanson
"""
"""
from sentence_transformers import SentenceTransformer
import chromadb
from transformers import AutoModelForCausalLM, AutoTokenizer

# ----------------------
# 1. 向量数据库构建（Chroma）
# ----------------------
# 初始化嵌入模型（用于将文本转换为向量）
embedding_model = SentenceTransformer("paraphrase-multilingual-MiniLM-L12-v2")

# 初始化 Chroma 客户端
chroma_client = chromadb.PersistentClient(path="rag_chroma_db")  # 数据持久化到本地
collection = chroma_client.get_collection(name="rag_demo")

# 添加数据到向量数据库
# nothings

# 加载 Qwen2-1.5B 模型和分词器
model_name = r"E:\soft\model\qwen\Qwen\Qwen2___5-1___5B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")

# ----------------------
#  检索增强生成（RAG）查询结果，score_threshold 最低分数值
# ----------------------
def rag_query(query_text, top_k=2, score_threshold=0.7):
    """RAG 流程：检索 + 生成 过滤低相似度结果"""
    # Step 1: 检索相关文档
    query_embedding = embedding_model.encode(query_text).tolist()
    results = collection.query(
        query_embeddings=[query_embedding],
        n_results=top_k
    )

    # Step 2: 过滤低相似度结果（Chroma 返回的是距离，需转换为分数）
    filtered_docs = []
    for doc, metadata, distance in zip(results["documents"][0],
                                       results["metadatas"][0],
                                       results["distances"][0]):
        similarity_score = 1 - distance  # 余弦相似度 = 1 - 欧氏距离（需确认 Chroma 的距离类型）
        if similarity_score >= score_threshold:
            filtered_docs.append({
                "text": doc,
                "score": similarity_score,
                "output": metadata["答案"]  # 保留原始输出
            })
            return metadata["答案"].split("答案：")[-1].strip()

# 如果未找到足够相似的文档，则返回默认答案 有大模型回答
    if not filtered_docs:
        print("未找到足够相关的上下文")
       #  return "未找到足够相关的上下文" # 返回默认答案 转交给大模型回答

        # 拼接检索到的上下文
        context = query_text
        # Step 2: 用 Qwen2 生成回答
        prompt = f"""基于以下上下文回答问题：
            {context}
            问题：{query_text}
            答案："""

        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        outputs = model.generate(**inputs, max_new_tokens=200)
        answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
        #return answer.split("答案：")[-1].strip()  # 提取生成部分
        return "以下回答来自 qwen2.5:"+answer.split("答案：")[-1].strip()


# ----------------------
# 4. 测试案例
# ----------------------
if __name__ == "__main__":
    query = "什么东西别人请你吃，但你自己还是要付钱？"
    answer = rag_query(query)
    print(f"问题：{query}\n答案：{answer}")
    # 示例查询
    query = "mybatis 是什么框架？"
    answer = rag_query(query)
    print(f"问题：{query}\n答案：{answer}")


